import pandas as pd
from sklearn.model_selection import train_test_split
import os

# 读取 CSV 文件
file_path = 'E:\人工智能\资料\专高六\数据清洗\数据\（分类）train_cleaned.csv'
df = pd.read_csv(file_path)

# 目标文件夹
target_folder = 'E:\人工智能\资料\专高六\数据清洗\Bert-Chinese-Text-Classification-Pytorch\THUCNews'
if not os.path.exists(target_folder):
    os.makedirs(target_folder)

# 划分训练集、验证集和测试集，比例为 8:1:1
train_df, temp_df = train_test_split(df, test_size=0.2, random_state=42)
dev_df, test_df = train_test_split(temp_df, test_size=0.5, random_state=42)

# 定义保存函数
def save_to_txt(df, file_path):
    with open(file_path, 'w', encoding='utf-8') as f:
        for index, row in df.iterrows():
            sentence = row['sentence']
            label = int(row['label'])
            f.write(f'{sentence}\t{label}\n')

# 保存数据到对应的 txt 文件
save_to_txt(train_df, os.path.join(target_folder, 'train.txt'))
save_to_txt(dev_df, os.path.join(target_folder, 'dev.txt'))
save_to_txt(test_df, os.path.join(target_folder, 'test.txt'))

print("数据已成功保存到指定文件夹的 txt 文件中。")